import cv2
import numpy as np


def cal_conv(src_area, core, core_height, core_width):
    """
    手动计算卷积，输入为原图像的部分区域（与卷积核等大）
    :param src_area: ndarray格式，deltaH*deltaW
    :param core: ndarray格式，deltaH*deltaW
    :param core_height: 卷积核高deltaH
    :param core_width: 卷积核宽deltaW
    :return:
    """
    res = 0
    for i in range(0, core_height):
        for j in range(0, core_width):
            res += src_area[i, j] * core[i, j]
    return res


def edgeEnhancement(src: np.ndarray, core: np.ndarray, alpha=0.05):
    """
    边缘增强函数
    :param src: 数据类型numpy ndarray，范围0-255 uint8，图像形状H*W*3，顺序RGB，或H*W灰度图
    :param core: 数据类型numpy ndarray，卷积核大小需要为奇数，形状core_height*core_width
    :param alpha: 边缘以多大比例叠加到原图上
    :return:
    """
    core_height, core_width = core.shape
    s = (core_height - 1) // 2
    t = (core_width - 1) // 2
    res = np.zeros_like(src)
    # 对单通道图和三通道图分别处理
    if len(src.shape) == 3:
        src_height, src_width, src_channel = src.shape
        src_pad = np.pad(src, ((s, s), (t, t), (0, 0)))
        for h in range(1, src_height + 1):
            for w in range(1, src_width + 1):
                for c in range(0, src_channel):
                    src_area = src_pad[h - s:h + s + 1, w - t:w + t + 1, c]
                    res[h - 1, w - 1, c] = cal_conv(src_area, core, core_height, core_width)
    else:
        src_height, src_width = src.shape
        src_pad = np.pad(src, ((s, s), (t, t)))
        for h in range(1, src_height + 1):
            for w in range(1, src_width + 1):
                src_area = src_pad[h - s:h + s + 1, w - t:w + t + 1]
                res[h - 1, w - 1] = cal_conv(src_area, core, core_height, core_width)

    res = np.clip(res, 0, 255)
    res = (alpha * res + (1 - alpha) * src).astype(np.uint8)
    return res


if __name__ == '__main__':
    # 读入图像，注意BGR
    src = cv2.imread('4g.jpg')  # H*W*C
    src = src[:, :, 2:]

    print(src.shape)
    src_height, src_width, src_channel = src.shape
    # core = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
    # core = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
    core = np.array([[1, 1, 1], [1, -8, 1], [1, 1, 1]])
    # core = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]])
    # core = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
    # core = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
    # core = np.array([[-1, 0], [0, 1]])
    # core = np.array([[0, -1], [1, 0]])

    # print(core)
    core_height, core_width = core.shape
    s = (core_height - 1) // 2
    t = (core_width - 1) // 2

    src_pad = np.pad(src, ((s, s), (t, t), (0, 0)))
    # print(src_pad[:,:,0])
    res = np.zeros_like(src)
    for h in range(1, src_height + 1):
        for w in range(1, src_width + 1):
            for c in range(0, src_channel):
                src_area = src_pad[h - s:h + s + 1, w - t:w + t + 1, c]
                # print(src_area.shape)
                # print(src_area)
                res[h - 1, w - 1, c] = cal_conv(src_area, core, core_height, core_width)
    res = np.clip(res, 0, 255)
    # print(res[:,:,0])
    # cv2.imshow("0", res)
    cv2.imshow("0", (0.05 * res + 0.95 * src).astype(np.uint8))
    cv2.waitKey(0)
